from tvtk.api import tvtk
from mayavi.scripts import mayavi2

from tvtk.array_handler import array2vtk

import numpy
from meshpy_examples import shocktube2D
import grid

def get_gmsh_mesh():
    import gmsh_read as gmsh
    g = gmsh.Gmsh2D('gmsh_examples/wedge.msh')
    
    points = g.points
    elements = g.elements
    quad_type = tvtk.Quad().cell_type

    ug = tvtk.UnstructuredGrid(points=points)
    ug.set_cells(quad_type, elements)

    return ug, g

def get_mesh():
    grid = shocktube2D()

    _mesh = grid._mesh

    points = numpy.zeros( shape=(grid.npoints, 3) )
    points[:, 0] = grid.points[:, 0]
    points[:, 1] = grid.points[:, 1]

    triangles = grid.elements
    triangle_type = tvtk.Triangle().cell_type

    ug = tvtk.UnstructuredGrid(points=points)
    ug.set_cells(triangle_type, triangles)
    
    return ug, grid

#ug, g = get_mesh()
ug, g = get_gmsh_mesh()

npoints = len(ug.points)
velocity = numpy.random.randn(npoints, 3)
temp = numpy.random.random(npoints)

ncells = g.nelements
nghost = g.nghost

cellid = numpy.zeros( ncells )
cellid[-nghost:] = 1

cellid = tvtk.to_tvtk( array2vtk(cellid) ); cellid.name = "id"

ncells = ug.number_of_cells
rho = tvtk.to_tvtk( array2vtk(numpy.ones(ncells)) ); rho.name = "density"
p = tvtk.to_tvtk( array2vtk(numpy.ones(ncells) * 2) ); p.name = "pressure"

ug.cell_data.add_array(p)
ug.cell_data.add_array(rho)
ug.cell_data.add_array(cellid)

ug.cell_data.set_active_scalars("id")

ug.point_data.vectors = velocity
ug.point_data.vectors.name = "velocity"

# @mayavi2.standalone
# def view():
#     from mayavi.sources.vtk_data_source import VTKDataSource
#     from mayavi.modules.outline import Outline
#     from mayavi.modules.surface import Surface
#     from mayavi.modules.vectors import Vectors

#     mayavi.new_scene()
#     src = VTKDataSource(data = ug)
#     mayavi.add_source(src)

#     mayavi.add_module(Outline())

#     surface = Surface()

#     mayavi.add_module(surface)
#     surface.actor.property.representation = "surface"
#     surface.actor.property.shading = "true"
#     surface.actor.property.edge_visibility = True
#     surface.actor.property.line_width = 0.5

#     engine = mayavi.engine
#     scene = engine.scenes[0]
#     vtk_src = scene.children[0]
#     mm = vtk_src.children[0]
#     modules = mm.children
#     surface = modules[1]

#     # surface.actor.property.linewidth = 2.0
#     # surface.actor.property.representation = "wireframe"
#     # vtk_src.point_scalars_name = ''

if __name__ == "__main__":
    view()
